import pandas as pd
import numpy as np
from sklearn.datasets import load_iris
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 加载Iris数据集
iris = load_iris()
df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
df['species'] = iris.target

# 数据标准化
scaler = StandardScaler()
df_scaled = scaler.fit_transform(df[df.columns[:-1]])

# 应用K-means算法
num_clusters = 3  # Iris数据集中有三种不同的花种类
kmeans = KMeans(n_clusters=num_clusters, random_state=42)
df['cluster'] = kmeans.fit_predict(df_scaled)

# 可视化聚类结果
fig = plt.figure(figsize=(10, 7))
ax = fig.add_subplot(111, projection='3d')

# 使用前三个特征进行三维可视化
scatter = ax.scatter(df.iloc[:, 0], df.iloc[:, 1], df.iloc[:, 2],
                     c=df['cluster'], cmap='viridis', s=50)

# 添加颜色条
legend1 = ax.legend(*scatter.legend_elements(), title="Clusters")
ax.add_artist(legend1)

# 设置轴标签
ax.set_xlabel('Sepal Length')
ax.set_ylabel('Sepal Width')
ax.set_zlabel('Petal Length')
plt.title('K-means Clustering on Iris Dataset')
plt.show()

# 输出每个簇中样本的数量
print("Cluster distribution:")
print(df['cluster'].value_counts())

# 比较实际类别与聚类结果（仅用于评估）
print("\nComparison with actual species:")
print(pd.crosstab(df['species'], df['cluster']))